"""
使用向量存储作为检索器
向量存储检索器是一个使用向量存储来检索文档的检索器。它是一个轻量级的包装器，围绕向量存储类构建，使其符合检索器接口。 它使用向量存储实现的搜索方法，如相似性搜索和MMR，来查询向量存储中的文本。
"""
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import CharacterTextSplitter

from models import get_ollama_embeddings_client

file_path = "../data/document/公司管理制度.txt"
loader = TextLoader(file_path, encoding="utf-8")
documents = loader.load()

text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0)
texts = text_splitter.split_documents(documents)

embeddings_client = get_ollama_embeddings_client()
vector_store = Chroma.from_documents(texts, embeddings_client)
retriever = vector_store.as_retriever()
# 最大边际相关性检索
# retriever = vector_store.as_retriever(search_type="mmr")

# 传递搜索参数
# 指定前K
# retriever = vector_store.as_retriever(search_kwargs={"k": 1})
# 相似性得分阈值检索，仅返回得分高于该阈值的文档。
# retriever = vector_store.as_retriever(search_type="similarity_score_threshold", search_kwargs={"score_threshold": 0.5})

docs = retriever.invoke("晋升")
print(docs)
